# -*- coding: utf-8 -*-
# time: 2025/3/29 10:55
# file: ch01.py
# author: hanson
"""
Embedding Models：chromadb 内置的 all-MiniLM-L6-v2

"""
import chromadb.config
import ollama
from flask import Flask, request, jsonify

app = Flask(__name__)
# 初始化 ChromaDB 客户端和集合
client = chromadb.PersistentClient(path="data/chromadb_data")
collection = client.get_or_create_collection(name="docs")
@app.route("/add_document", methods=["POST"])
def add_document():
    # 从请求中获取文档
    data = request.json
    documents = data.get("documents", [])

    # 将每个文档存储在向量嵌入数据库中
    for i, d in enumerate(documents):
        collection.upsert(documents=d, ids=[str(i)])

    return jsonify({"message": "Documents added successfully!"}), 201

@app.route("/query", methods=["POST"])
def query():
    data = request.json
    prompt = data.get("prompt", "")
    # 查询 ChromaDB
    results = collection.query(
        query_texts=[prompt],
        n_results=2,
    )
    print(data,prompt)
    if results["documents"]:
        print(results["documents"])
        data = ollama.generate( model="deepseek-r1:1.5b",
                                # system_prompt=results["documents"][0][0],
                                 prompt=f"根据这段文字：{data}。回答这个问题：{prompt}",
                                 # temperature=0.7,
                                 # max_tokens=512,
                                 # top_p=0.95,
                                 # top_k=40,
                                 # repeat_penalty=1.1,
                                 # repeat_last_n=64,
                                  )
        print(data.response)
        return data.response, 200
    else:
        return jsonify({"message": "No documents found."}), 404

if __name__ == "__main__":
    app.run(debug=True)

